#!/bin/bash

export PYTHONPATH=../tnet/python:../torn

die() {
  echo 1>&2 "FATAL: $@"
  exit 1
}

# Run unit tests (with current PYTHONPATH)
unit() {
  local test_path=$PWD/$1
  shift
  $test_path "$@"
}

# Run with PYTHONPATH
xmap() {
  ./xmap.py "$@"
}


readonly POLY_DIR=~/hg/polyweb

# TODO: Source tin-info.sh and call tin_info here.
build-tin() {
  local main=${1:-xmap}
  # Need '.' in PYTHONPATH to import xmap.py and build tin file
  export PYTHONPATH=.:$PYTHONPATH
  $POLY_DIR/tools/py_imports.py $main \
     | grep -v /usr/lib/python  \
     | $POLY_DIR/tools/tin.py
}

# Build all
build-all() {
  set -o errexit
  build-tin xmap
  build-tin sleep_
  build-tin exec_
  build-tin chunks
  ls -al *.tin
}


#
# Helpers
#

to-json() {
  ../tnet/python/tnet.py to-json
}

ints() {
  ./chunks.py -t '#' "$@"
}

floats() {
  ./chunks.py -t '^' "$@"
}


sleep() {
  time ints 1 2 | ./sleep_.py | ../tnet/python/tnet.py to-json
}


#
# Xmap demos
#

xmap-echo() {
  time ints 1 2 | ./xmap.py ./sleep_.py "$@"
}

xmap-broken() {
  time ints 1 2 | ./break.py | ./xmap.py ./sleep_.py "$@"
}

xmap-2() {
  echo 'Should take 1 second'
  time ints 1 1 | ./xmap.py --max-procs 2 ./sleep_.py "$@"
}

xmap-3() {
  echo 'Should take 2 seconds'
  time ints 1 1 1 | xmap --max-procs 2 ./sleep_.py "$@"
}

xmap-4() {
  echo 'Should only start 3 processes!'
  time ints 1 1 1 | ./xmap.py --max-procs 20 ./sleep_.py
}

xmap-uneven() {
  # 1 10 10 1
  # this takes 11 seconds.  But if you to '1 10 1 10', it takes 20 seconds.  I
  # don't think there's a way to know this in general, unless xmap has a concept
  # of extracting request "cost".
  time ints 1 10 10 1 | ./xmap.py --max-procs 2 ./sleep_.py
}
xmap-uneven2() {
  # 1 10 1 1 10 -- try to balance it evenly?
  time ints 1 10 1 1 10 | ./xmap.py --max-procs 2 ./sleep_.py
}

xmap-inf() {
  time floats 0.2 | repeat 100000 \
    | ./xmap.py --log-dir sleep -- ./sleep_.py "$@" \
    | cat > sleep.txt
}

xmap-output() {
  echo 'Testing output'
  time ints 1 1 1 | ./xmap.py --max-procs 2 ./sleep_.py | to-json
}

# TODO: Make these 2 test cases pass!
xmap-infile() {
  ./xmap.py ./sleep_.py < in.txt
}

xmap-outfile() {
  time ints 1 | ./xmap.py ./sleep_.py > out.txt
}

# Just use exec_
exec-curl() {
  # chunks turns lines into strings, and then we have a template in exec_.py
  time cat testdata/urls.txt | ./chunks.py | ./exec_.py -- curl --include _
}

# Parallelize with xmap exec_
par-exec-curl() {
  time cat testdata/urls.txt | ./chunks.py \
    | ./xmap.py --log-dir tmp -P 3 -- ./exec_.py -- curl --include _ \
    | cat > pages.tnet
}

to-chunks() {
  # Input: byte stream of space separated strings
  # Output: Array of strings in TNET format
  sed 's/ /\n/g' | ./chunks.py --lines -t ',' | ./chunks.py -t ']'
}

single() {
  to-chunks | ./exec_.py
}

repeat() {
  # Repeat input stream N times

  cat > /tmp/repeat
  local n=$1
  test -n "$n" || { echo "N required"; exit 1; }
  for i in $(seq $n); do
    cat /tmp/repeat
  done
}

# Input: a raw command
# Output: TNET chunks suitable for exec
string2exec() {
  local n=${1:-10}
  to-chunks | repeat $n
}

par() {
  # Run a command 10 times with default parallelism
  # stdin: space-separated strings
  string2exec | ./xmap.py --log-dir logs "$@" -- ./exec_.py
}

#
# Local
#

recursive-local() {
  ./xmap.py --argv-strings testdata/procs.txt
}

one-six() {
  ints 1 2 3 4 5 6
}

one-100 () {
  floats 1.0 | repeat 100
}

# 0.11 to make the seconds fractional
tenth-1000 () {
  floats 0.11 | repeat 1000
}

recursive-local-ui() {
  ./xmap.py --argv-strings testdata/muxed.txt --child-demux output,status \
    | cat >/dev/null
}

recursive-remote() {
  ./xmap.py --argv-strings testdata/remote.txt
}

xmap-outmux() {
  time ints 1 1 1 | ./xmap.py --output @1 --status-out @2 ./sleep_.py "$@" \
    | cat > outmux.tnet
  echo
  cat outmux.tnet
  echo
  cat outmux.tnet | to-json
}

xmap-outmux2() {
  time floats 0.3 | repeat 10 \
    | ./xmap.py --max-procs 2 --log-dir tmp \
      --output @1 --status-out @2 \
      ./sleep_.py "$@" #\
    #| cat > outmux2.tnet
}

# 
# REMOTE
# 

servers() {
  local filename=${1:-testdata/servers.txt}
  cat $filename | grep -v '#' 
}

make-payload() {
  local manifest="$1"
  test -n "$manifest" || die "Manifest required"

  set -o errexit

  # Create a dir
  rm -rf payload
  mkdir payload
  { cat $1;
    echo xmap.tin;
    echo sleep_.tin;
    echo exec_.tin
  } | xargs --verbose -I {} -- cp {} payload

  # Create a tar file from the dir, preserving permissions
  tar --create --preserve-permissions --verbose --file payload.tar payload/*

  echo
  echo CONTENTS
  echo

  tar --list --file payload.tar
}

# Use xmap to copy payload.tar to a list of servers
# TODO: Need quoting to do tar piping trick
copy-payload() {
  servers "$@" | row '$1:/tmp/' | ./chunks.py --lines \
    | ./xmap.py --log-dir copy -- ./exec_.py scp payload.tar _
}

copy-payload-seq() {
  servers "$@" \
    | xargs -I {} --verbose -- \
      sh -c 'cat payload.tar | ssh {} tar xpf - -C /tmp'
}

make-out-seq() {
  servers "$@" \
    | xargs -I {} --verbose -- \
      ssh {} mkdir /tmp/payload/fargo_out
}

# OLD: copy specific files
copy() {
  servers "$@" | row '$1:/tmp/' | ./chunks.py --lines \
    | ./xmap.py --log-dir copy -- ./exec_.py scp xmap.tin sleep_.tin exec_.tin _
  servers "$@" | ./chunks.py --lines \
    | ./xmap.py --log-dir copy -- ./exec_.py \
      ssh _ chmod 755 /tmp/xmap.tin /tmp/sleep_.tin /tmp/exec_.tin
  cp --verbose xmap.tin sleep_.tin exec_.tin /tmp
}

# How to do multiple servesr and multiple tasks?
# Really xmap.py has to start up ssh itself... that is most consistent.  Then
# flow control will make it so it takes up the most tasks.
remote()  {
  host=chubot.com
  time ints 1 1 1 \
    | ssh $host /tmp/xmap.tin --no-term -- /tmp/sleep_.py
}

make-remote-argv() {
  local hosts=$1
  if test -n "$hosts"; then
    hosts="$(cat $hosts)"
  else
    hosts="192.168.0.2" # chubot@chubot.org"
  fi
  echo '# remote-argv.txt' > remote-argv.txt
  for host in $hosts; do
    echo "ssh $host /tmp/xmap.tin --output @1 --status-out @2 --no-term --log-dir /tmp -- /tmp/sleep_.tin" >> remote-argv.txt
  done
  # Local copy too
  echo "/tmp/xmap.tin --output @1 --status-out @2 --no-term --log-dir /tmp -- /tmp/sleep_.tin" >> remote-argv.txt
}

make-local-argv() {
  procs="1 2"  # TODO: name these
  echo '# local-argv.txt' > local-argv.txt
  for proc in $procs; do
    echo "./xmap.tin --output @1 --status-out @2 --no-term --log-dir /tmp -- ./sleep_.tin" >> local-argv.txt
  done
}

remote-recursive() {
  local argv=${1:-remote-argv.txt}
  ./xmap.tin --argv-strings $argv --log-dir tmp --child-demux output,status \
    | cat >/dev/null
}

make-exec-argv() {
  local hosts=$1
  if test -n "$hosts"; then
    hosts="$(cat $hosts)"
  else
    hosts="192.168.0.2" # chubot@chubot.org"
  fi
  echo '# remote-exec.txt' > remote-exec.txt
  for host in $hosts; do
    echo "ssh $host /tmp/payload/xmap.tin --output @1 --status-out @2 --no-term --log-dir /tmp -- /tmp/payload/exec_.tin -d /tmp/payload" >> remote-exec.txt
  done
  # Local copy too
  #echo "/tmp/payload/xmap.tin --output @1 --status-out @2 --no-term --log-dir /tmp -- /tmp/payload/exec_.tin -d /tmp/payload" >> remote-exec.txt
}


"$@"
